import cv2
import numpy as np

from .base import Preprocessor


class Gray2RGB(Preprocessor):
    def __init__(self):
        super().__init__()

    def process(self, images):
        """
        使用 cv2 把灰度图 (V, H, W, 1) 转为 RGB (V, H, W, 3)
        """
        if images.ndim != 4 or images.shape[-1] != 1:
            raise ValueError(f"期望输入 (V, H, W, 1)，实际得到 {images.shape}")

        rgb_frames = []
        for idx, img in enumerate(images):
            gray = img.squeeze(-1)  # (H, W)
            # 新版 OpenCV（>=4.5.2）：
            rgb = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
            # 若你的版本只支持 COLOR_GRAY2BGR，可改为：
            # bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
            # rgb = cv2.cvtColor(bgr,  cv2.COLOR_BGR2RGB)
            rgb_frames.append(rgb)

        return np.stack(rgb_frames, axis=0).astype(images.dtype)
